package com.milageincrease.api;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.Connection;
import org.apache.hadoop.hbase.client.ConnectionFactory;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.Table;
import org.apache.hadoop.hbase.util.Bytes;

import java.io.IOException;
import java.util.ArrayList;

/**
 * HBase的写入API
 *
 * @author Hanzhuo
 */
public class HBaseWriteTools {
    private Table table;
    private final boolean initTableName;
    private static final int ROW_KEY_LENGTH = 13;
    private static final Configuration CONF = HBaseConfiguration.create();
    public static Connection CONN;

    static {
        try {
            CONN = ConnectionFactory.createConnection(CONF);
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * 无参构造方法，未指定TableName <br/>
     * 需要调用put(String tableName, JSONObject data)进行数据写入 <br/>
     * 仅适用于插入少量数据
     */
    @Deprecated
    public HBaseWriteTools(){
        this.initTableName = false;
    }

    /**
     * 推荐，初始化时即打开HBase链接，需调用close()方法关闭链接<br/>
     * 调用put(JSONObject data)进行数据写入，<br/>
     * @param tableName 初始化表名，并保持连接，方便持续插入数据
     */
    public HBaseWriteTools(String tableName) throws IOException {
        this.initTableName = true;
        this.table = CONN.getTable(TableName.valueOf(tableName));
    }

    /**
     * 适合连续多次，每次少量且间隔长的插入
     * @param data 数据，JSON格式
     * @throws HBaseTableNameInitException 构造本对象时未指定HBase表名
     */
    public void put(JSONObject data) throws HBaseTableNameInitException, RowKeyFormatException, IOException {
        if (!this.initTableName) {
            String message = "未指定HBase表名，请调用 put(String tableName, JSONObject data) 方法或使用HBaseWriteAPI()构造方法";
            throw new HBaseTableNameInitException(message);
        }
        String rowKey = String.valueOf(data.getString("rowkey"));
        verifyRowKey(rowKey);
        Put put = parser(data.getJSONObject("data"), rowKey);
        this.table.put(put);
    }

    /**
     * 适合批量、大量插入数据，参数传入为一些JSONObject封装而成的的JSONArray
     * @param dataList JSON列表
     * @throws HBaseTableNameInitException 构造本对象时未指定HBase表名
     * @throws RowKeyFormatException RowKey格式错误
     * @throws IOException .
     */
    public void put(JSONArray dataList) throws HBaseTableNameInitException, RowKeyFormatException, IOException {
        if (!this.initTableName) {
            String message = "未指定HBase表名，请调用 put(String tableName, JSONObject data) 方法或使用HBaseWriteAPI()构造方法";
            throw new HBaseTableNameInitException(message);
        }
        ArrayList<Put> puts = new ArrayList<>();
        for (int i = 0; i < dataList.size(); i++) {
            String rowKey = String.valueOf(dataList.getJSONObject(i).getString("rowkey"));
            verifyRowKey(rowKey);
            Put put = parser(dataList.getJSONObject(i).getJSONObject("data"), rowKey);
            puts.add(put);
        }
        this.table.put(puts);
    }

    /**
     * 数据量较小，使用该方法，方便 <br/>
     * Warning: 在插入数据量较大时，请勿使用该方法，应该使用pipline，否则将导致资源耗尽
     */
    @Deprecated
    public void put(String tableName, JSONObject data) throws IOException, HBaseTableNameInitException, RowKeyFormatException {
        if (this.initTableName){
            String message = "已经指定HBase表名，请调用 put(JSONObject data) 方法";
            throw new HBaseTableNameInitException(message);
        }
        this.table = CONN.getTable(TableName.valueOf(tableName));
        String rowKey = String.valueOf(data.getString("rowkey"));
        verifyRowKey(rowKey);
        Put put = parser(data.getJSONObject("data"), rowKey);
        this.table.put(put);
        this.table.close();
    }

    /**
     * 关闭HBaseTable的连接
     * @throws IOException 关闭失败
     */
    public void close() throws IOException {
        this.table.close();
    }

    /**
     * 内部对象，解析JSON，返回Put对象
     * @param data JSON数据
     * @param rowKey HBase rowKey
     * @return Put Object
     */
    private Put parser(JSONObject data, String rowKey){
        Put put = new Put(Bytes.toBytes(rowKey));
        for (String cfKey: data.keySet()){
            JSONObject cf = data.getJSONObject(cfKey);
            for (String cKey: cf.keySet()){
                put.addColumn(cfKey.getBytes(), cKey.getBytes(), String.valueOf(cf.get(cKey)).getBytes());
            }
        }
        return put;
    }

    /**
     * 验证RowKey是否符合规范
     */
    private void verifyRowKey(String rowKey) throws RowKeyFormatException {
        if (rowKey.length() != ROW_KEY_LENGTH) {
            String message = String.format("rowKey: %s 未通过校验，请检查格式", rowKey);
            throw new RowKeyFormatException(message);
        }
    }

    /*
      获取插入数据的JSON格式<br/>
      {
          "rowkey": "random.nextInt(100):dataStationId:TimeStamp" // rowkey的random.nextInt(100)为两位随机整数，如果数字小于10需要用0补齐
          "data": {
              "ColumnFamilyName1": {
                  "ColumnName1": "value1",
                  "ColumnName2": "value2",
                  ...
              },
              "ColumnFamilyName2": {
                  "ColumnName3": "value4",
                  ...
              },
              ...
          }
      }
     */
}
